package com.zxjbyte.yiyi.module.plugin.sse.session;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.zxjbyte.yiyi.framework.common.cache.caffeine.impl.CaffeineCache;
import com.zxjbyte.yiyi.framework.common.constant.HeaderMsg;
import com.zxjbyte.yiyi.framework.common.domain.Kv;
import com.zxjbyte.yiyi.framework.common.domain.message.StreamMessage;
import com.zxjbyte.yiyi.module.plugin.sse.processor.HeartbeatProcessor;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * SSE会话
 *
 * @Author zhangxingjia
 * @Date 2024/7/12 15:21
 * @Version: 1.0
 */
@Slf4j
public class SseSession {

    // 默认的客户端
    private static final String DEFAULT_CLIENT_ID = "yiyi-boot";
    // 当前用户链接数
    private static final CaffeineCache<String, AtomicInteger> count = CaffeineCache.newCaffeineCache();
    // 会话缓存
    @Getter
    private static final CaffeineCache<String, Map<String, Kv>> sseSessionCache = CaffeineCache.newCaffeineCache();

    /**
     * 推送消息到客户端
     *
     * @param clientId 客户端标识
     * @param userId 用户id
     * @param message 推送消息
     */
    public static boolean sendMessage(String clientId, String userId, StreamMessage message){
        clientId = getClientId(clientId);
        if(message == null){
            log.warn("SSE send failed message is empty, clientId:{}, userId:{}", clientId, userId);
            return false;
        }
        Map<String, Kv> stringKvMap = sseSessionCache.get(clientId);
        if(stringKvMap == null){
            log.warn("SSE send message failed, clientId:{}, userId:{}", clientId, userId);
            return false;
        }
        Kv kv = stringKvMap.get(userId);
        if(kv == null){
            log.warn("SSE send message failed, clientId:{}, userId:{}", clientId, userId);
            return false;
        }
        SseEmitter sseEmitter = (SseEmitter) kv.get(SseSessionKeyEnum.EMITTER.name());
        return executeSend(sseEmitter, message);
    }

    public static boolean sendMessage(String clientId, String userId, String message){
        return sendMessage(clientId, userId, StreamMessage.of(HeaderMsg.SEND_MESSAGE, message));
    }

    public static boolean sendMessage(String userId, StreamMessage message){
        return sendMessage(null, userId, message);
    }

    /**
     * 广播消息到客户端
     *
     * @param clientId 客户端标识
     * @param message 消息
     */
    public static void broadcastMessage(String clientId, StreamMessage message){
        clientId = getClientId(clientId);
        if(message == null){
            log.warn("SSE broadcast failed message is empty, clientId:{}", clientId);
            return;
        }
        Map<String, Kv> stringKvMap = sseSessionCache.get(clientId);
        if(CollUtil.isNotEmpty(stringKvMap)){
            for(Map.Entry<String, Kv> entry : stringKvMap.entrySet()){
                SseEmitter sseEmitter = (SseEmitter) entry.getValue().get(SseSessionKeyEnum.EMITTER.name());
                executeSend(sseEmitter, message);
            }
        }
    }

    public static void broadcastMessage(StreamMessage message){
        broadcastMessage(null, message);
    }

    /**
     * 广播客户端在线用户数
     * @param clientId
     */
    public static void broadcastOnlineUsersNumber(String clientId){
        int onlineUsersNumber = getOnlineUsersNumber(clientId);
        broadcastMessage(clientId, StreamMessage.of(HeaderMsg.ONLINE_USERS_NUMBER, onlineUsersNumber));
    }

    /**
     * 广播客户端在线用户数
     */
    public static void broadcastOnlineUsersNumber(){
        broadcastOnlineUsersNumber(null);
    }

    /**
     * 获取在线用户数
     * @param clientId
     * @return
     */
    public static int getOnlineUsersNumber(String clientId){
        AtomicInteger atomicInteger = count.get(getClientId(clientId));
        return atomicInteger == null ? 0 : atomicInteger.get();
    }

    /**
     * 根据客户端id和用户id获取连接对象
     * @param clientId
     * @param userId
     * @return
     */
    public static SseEmitter getSseEmitter(String clientId, String userId){
        Map<String, Kv> stringKvMap = sseSessionCache.get(getClientId(clientId));
        if(CollUtil.isNotEmpty(stringKvMap)){
            Kv kv = stringKvMap.get(userId);
            if(kv != null){
                return (SseEmitter) kv.get(SseSessionKeyEnum.EMITTER.name());
            }
        }
        return null;
    }

    /**
     * 新增连接
     */
    public static void addConnection(String clientId, String userId, SseEmitter emitter, Boolean openHeartBeat){
        boolean isFirst = sseSessionCache.isEmpty();
        Map<String, Kv> stringKvMap = sseSessionCache.get(getClientId(clientId), k -> new ConcurrentHashMap<>());
        stringKvMap.put(userId, Kv.create().set(SseSessionKeyEnum.EMITTER.name(), emitter).setIgnoreNull(SseSessionKeyEnum.HEARTBEAT.name(), openHeartBeat));
        // 新增连接数 + 1
        getCount(clientId).getAndIncrement();
        // 第一个连接开启全局心跳
        if(isFirst){
            HeartbeatProcessor.broadcastHeartbeat();
        }
    }


    /**
     * 移除连接
     */
    public static void removeConnection(String clientId, String userId){
        clientId = getClientId(clientId);
        Map<String, Kv> stringKvMap = sseSessionCache.get(clientId);
        if(stringKvMap != null){
            stringKvMap.remove(userId);
            if(stringKvMap.isEmpty()){
                sseSessionCache.remove(clientId);
            }
        }
        // 连接数 - 1
        getCount(clientId).getAndDecrement();
    }


    private static boolean executeSend(SseEmitter sseEmitter, StreamMessage message){
        SseEmitter.SseEventBuilder sendData = SseEmitter.event().data(message, MediaType.APPLICATION_JSON);
        try {
            Objects.requireNonNull(sseEmitter).send(sendData);
        } catch (Exception e) {
            // 标记 SseEmitter 已完成，并附带一个错误, 触发 onCompletion 回调
            sseEmitter.completeWithError(e);
            return false;
        }
        return true;
    }


    private static AtomicInteger getCount(String clientId) {
        return count.get(getClientId(clientId), k -> new AtomicInteger(0));
    }

    /**
     * 这里的clientId 实际为保留的扩展字段，如果只是本系统使用，可以为空取默认
     * 当多客户端使用时可传clientId标识
     * @param clientId
     * @return
     */
    private static String getClientId(String clientId){
        if(StrUtil.isEmpty(clientId)){
            return DEFAULT_CLIENT_ID;
        }
        return clientId;
    }

    /**
     * 会话缓存key枚举
     */
    public enum SseSessionKeyEnum{
        // 通信
        EMITTER,
        // 心跳
        HEARTBEAT,
    }
}
